package top.continew.minithreadpool.concurrent;

import top.continew.minithreadpool.factory.ThreadFactory;
import top.continew.minithreadpool.policy.RejectedExecutionHandler;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 线程池类，实现 ThreadPoolExecutor 接口
 */
public class ThreadPool implements ThreadPoolExecutor {
//    private static final ThreadLocal<ThreadPool> currentThreadPool = new ThreadLocal<>();

    private final ThreadPoolState state = new ThreadPoolState();
    private volatile int corePoolSize;
    private volatile int maxPoolSize;
    private final long keepAliveTime;
    private final TimeUnit timeUnit;
    private final BlockingQueue<Runnable> workQueue;
    private final ThreadFactory threadFactory;
    private final RejectedExecutionHandler handler;

    private final Set<Worker> workers = new HashSet<>();
    private final ReentrantLock mainLock = new ReentrantLock();
    private final Condition termination = mainLock.newCondition();
    private final AtomicInteger workerCount = new AtomicInteger(0);
    private long completedTaskCount;

    /**
     * 创建线程池
     *
     * @param corePoolSize    核心线程数
     * @param maxPoolSize     最大线程数
     * @param keepAliveTime   线程空闲时间
     * @param unit            时间单位
     * @param workQueue       任务队列
     * @param threadFactory   线程工厂
     * @param handler         拒绝策略
     */
    public ThreadPool(int corePoolSize, int maxPoolSize, long keepAliveTime,
                      TimeUnit unit, BlockingQueue<Runnable> workQueue,
                      ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        if (corePoolSize < 0 || maxPoolSize <= 0 || maxPoolSize < corePoolSize || keepAliveTime < 0) {
            throw new IllegalArgumentException("Invalid thread pool parameters");
        }
        if (workQueue == null || threadFactory == null || handler == null) {
            throw new NullPointerException();
        }

        this.corePoolSize = corePoolSize;
        this.maxPoolSize = maxPoolSize;
        this.keepAliveTime = keepAliveTime;
        this.timeUnit = unit;
        this.workQueue = workQueue;
        this.threadFactory = threadFactory;
        this.handler = handler;
    }

//    static ThreadPool getCurrentThreadPool() {
//        return currentThreadPool.get();
//    }

    void runWorker(Worker worker) {
        Thread wt = Thread.currentThread();
        Runnable task = worker.firstTask;
        worker.firstTask = null;

        boolean completedAbruptly = true;
        try {
            while (task != null || (task = getTask()) != null) {
                // 不需要显式加锁

                // 检查线程池状态，确保线程未被中断
                if ((state.isStop() ||
                        (Thread.interrupted() && state.isStop())) &&
                        !wt.isInterrupted()) {
                    wt.interrupt();
                }

                try {
                    beforeExecute(wt, task);
                    Throwable thrown = null;
                    try {
                        task.run();
                    } catch (RuntimeException x) {
                        thrown = x;
                        throw x;
                    } catch (Error x) {
                        thrown = x;
                        throw x;
                    } catch (Throwable x) {
                        thrown = x;
                        throw new Error(x);
                    } finally {
                        afterExecute(task, thrown);
                    }
                } finally {
                    task = null;
                    worker.completedTasks++;
                }
            }
            completedAbruptly = false;
        } finally {
            processWorkerExit(worker, completedAbruptly);
        }
    }

    private Runnable getTask() {
        boolean timedOut = false; // 上次poll是否超时

        for (;;) {
            if (state.isStop() ||
                    (state.isShutdown() && workQueue.isEmpty())) {
                workerCount.decrementAndGet();
                return null;
            }

            int wc = workerCount.get();

            // 是否对worker进行超时控制
            boolean timed = wc > corePoolSize;

            if ((wc > maxPoolSize || (timed && timedOut))
                    && (wc > 1 || workQueue.isEmpty())) {
                if (workerCount.compareAndSet(wc, wc - 1)) {
                    return null;
                }
                continue;
            }

            try {
                Runnable r = timed ?
                        workQueue.poll(keepAliveTime, TimeUnit.NANOSECONDS) :
                        workQueue.take();
                if (r != null)
                    return r;
                timedOut = true;
            } catch (InterruptedException retry) {
                timedOut = false;
            }
        }
    }

    private void processWorkerExit(Worker worker, boolean completedAbruptly) {
        if (completedAbruptly) {
            workerCount.decrementAndGet();
        }

        mainLock.lock();
        try {
            completedTaskCount += worker.completedTasks;
            workers.remove(worker);
        } finally {
            mainLock.unlock();
        }

        tryTerminate();

        int c = workerCount.get();
        if (!completedAbruptly) {
            int min = workerCount.get() >= corePoolSize ? 0 : corePoolSize;
            if (min == 0 && !workQueue.isEmpty())
                min = 1;
            if (c >= min)
                return; // 不需要替换
        }
        addWorker(null, false);
    }

    private void tryTerminate() {
        for (;;) {
            if (state.isRunning() ||
                    state.isTerminated() ||
                    (state.isShutdown() && !workQueue.isEmpty())) {
                return;
            }

            if (workerCount.get() != 0) {
                interruptIdleWorkers(true);
                return;
            }

            mainLock.lock();
            try {
                if (state.getState() != ThreadPoolState.TERMINATED) {
                    state.setState(ThreadPoolState.TERMINATED);
                    termination.signalAll();
                    terminated();
                }
            } finally {
                mainLock.unlock();
            }
        }
    }

    private void interruptIdleWorkers(boolean onlyOne) {
        mainLock.lock();
        try {
            for (Worker worker : workers) {
                Thread t = worker.thread;
                if (!t.isInterrupted()) {
                    t.interrupt();
                }
                if (onlyOne) {
                    break;
                }
            }
        } finally {
            mainLock.unlock();
        }
    }

    /**
     * 提交任务到线程池执行
     * @param command
     */
    @Override
    public void execute(Runnable command) {
        // 验证任务是否为空
        if (command == null) throw new NullPointerException();

        // 线程池是否关闭
        if (state.isShutdown()) {
            handler.rejectedExecution(command, this);
            return;
        }

        // 工作线程数
        int wc = workerCount.get();
        // 是否需要创建工作线程
        if (wc < corePoolSize) {
            if (addWorker(command, true)) return;
            wc = workerCount.get();
        }

        // 尝试将任务加入队列，如果失败则尝试创建非核心线程，如果都失败则拒绝任务
        if (state.isRunning() && workQueue.offer(command)) {
            int recheck = workerCount.get();
            if (!state.isRunning() && remove(command)) {
                handler.rejectedExecution(command, this);
            } else if (recheck == 0) {
                addWorker(null, false);
            }
        } else if (!addWorker(command, false)) {
            handler.rejectedExecution(command, this);
        }
    }

    /**
     * 创建工作线程
     * @param firstTask
     * @param core
     * @return
     */
    private boolean addWorker(Runnable firstTask, boolean core) {
        // 无限循环重试创建工作线程
        retry:
        for (;;) {
            // 检查线程池状态
            if (state.isShutdown() &&
                    !(state.getState() == ThreadPoolState.SHUTDOWN && firstTask == null)) {
                return false;
            }

            // 检查工作线程数是否超过限制
            int wc = workerCount.get();
            int limit = core ? corePoolSize : maxPoolSize;
            if (wc >= limit) {
                return false;
            }

            // 原子性地增加工作线程计数
            if (workerCount.compareAndSet(wc, wc + 1)) {
                break retry;
            }
        }

        boolean workerStarted = false;
        boolean workerAdded = false;
        Worker w = null;
        try {
            // w = new Worker(firstTask, threadFactory);
            // 把worker添加到工作线程集合中
            w = new Worker(this, firstTask, threadFactory);
            final Thread t = w.thread;
            if (t != null) {
                mainLock.lock();
                try {
                    // 再次检查线程池状态
                    if (state.getState() < ThreadPoolState.SHUTDOWN ||
                            (state.getState() == ThreadPoolState.SHUTDOWN && firstTask == null)) {
                        if (t.isAlive()) {
                            throw new IllegalThreadStateException();
                        }
                        workers.add(w);
                        workerAdded = true;
                    }
                } finally {
                    mainLock.unlock();
                }

                if (workerAdded) {
                    t.start();
                    workerStarted = true;
                }
            }
        } finally {
            if (!workerStarted) {
                addWorkerFailed(w);
            }
        }
        return workerStarted;
    }

    private void addWorkerFailed(Worker w) {
        mainLock.lock();
        try {
            if (w != null) {
                workers.remove(w);
            }
            workerCount.decrementAndGet();
            tryTerminate();
        } finally {
            mainLock.unlock();
        }
    }

    private boolean remove(Runnable task) {
        return workQueue.remove(task);
    }

    @Override
    public void shutdown() {
        mainLock.lock();
        try {
            if (state.getState() >= ThreadPoolState.SHUTDOWN) {
                return;
            }
            state.setState(ThreadPoolState.SHUTDOWN);
            interruptIdleWorkers(false);
            onShutdown();
        } finally {
            mainLock.unlock();
        }
        tryTerminate();
    }

    @Override
    public List<Runnable> shutdownNow() {
        List<Runnable> tasks;
        mainLock.lock();
        try {
            state.setState(ThreadPoolState.STOP);
            interruptWorkers();
            tasks = drainQueue();
            onShutdown();
        } finally {
            mainLock.unlock();
        }
        tryTerminate();
        return tasks;
    }

    private void interruptWorkers() {
        mainLock.lock();
        try {
            for (Worker worker : workers) {
                worker.interruptIfStarted();
            }
        } finally {
            mainLock.unlock();
        }
    }

    private List<Runnable> drainQueue() {
        List<Runnable> taskList = new ArrayList<>();
        workQueue.drainTo(taskList);
        if (!workQueue.isEmpty()) {
            for (Runnable r : workQueue.toArray(new Runnable[0])) {
                if (workQueue.remove(r)) {
                    taskList.add(r);
                }
            }
        }
        return taskList;
    }

    @Override
    public boolean isShutdown() {
        return state.isShutdown();
    }

    @Override
    public boolean isTerminated() {
        return state.isTerminated();
    }

    @Override
    public boolean awaitTermination(long timeout, TimeUnit unit)
            throws InterruptedException {
        long nanos = unit.toNanos(timeout);
        mainLock.lock();
        try {
            while (!state.isTerminated()) {
                if (nanos <= 0) {
                    return false;
                }
                nanos = termination.awaitNanos(nanos);
            }
            return true;
        } finally {
            mainLock.unlock();
        }
    }

    @Override
    public int getCorePoolSize() {
        return corePoolSize;
    }

    @Override
    public void setCorePoolSize(int corePoolSize) {
        if (corePoolSize < 0) {
            throw new IllegalArgumentException();
        }

        int delta = corePoolSize - this.corePoolSize;
        this.corePoolSize = corePoolSize;

        if (workerCount.get() > corePoolSize) {
            interruptIdleWorkers(false);
        } else if (delta > 0) {
            int k = Math.min(delta, workQueue.size());
            while (k-- > 0 && addWorker(null, true)) {
                if (workQueue.isEmpty()) {
                    break;
                }
            }
        }
    }

    @Override
    public int getMaximumPoolSize() {
        return maxPoolSize;
    }

    @Override
    public void setMaximumPoolSize(int maximumPoolSize) {
        if (maximumPoolSize <= 0 || maximumPoolSize < corePoolSize) {
            throw new IllegalArgumentException();
        }
        this.maxPoolSize = maximumPoolSize;

        if (workerCount.get() > maximumPoolSize) {
            interruptIdleWorkers(false);
        }
    }

    @Override
    public int getActiveCount() {
        mainLock.lock();
        try {
            int count = 0;
            for (Worker worker : workers) {
                if (worker.isLocked()) {
                    count++;
                }
            }
            return count;
        } finally {
            mainLock.unlock();
        }
    }

    @Override
    public BlockingQueue<Runnable> getQueue() {
        return workQueue;
    }

    @Override
    public long getCompletedTaskCount() {
        mainLock.lock();
        try {
            long count = completedTaskCount;
            for (Worker worker : workers) {
                count += worker.completedTasks;
            }
            return count;
        } finally {
            mainLock.unlock();
        }
    }

    @Override
    public int getPoolSize() {
        mainLock.lock();
        try {
            return workers.size();
        } finally {
            mainLock.unlock();
        }
    }

    // 钩子方法
    protected void beforeExecute(Thread t, Runnable r) {}
    protected void afterExecute(Runnable r, Throwable t) {}
    protected void terminated() {}
    protected void onShutdown() {}

    @Override
    public String toString() {
        return "ThreadPool[core=" + corePoolSize + ", max=" + maxPoolSize +
                ", active=" + getActiveCount() + ", queue=" + workQueue.size() + "]";
    }

    public int getQueueCapacity() {
        if (workQueue instanceof LinkedBlockingQueue) {
            return ((LinkedBlockingQueue<Runnable>) workQueue).remainingCapacity() + workQueue.size();
        }
        return -1; // 未知容量
    }
}
